/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ola_security.h"

#include <stddef.h>

#include "ola_base_utils.h"
#include "ola_def.h"
#include "ola_error.h"
#include "ola_pal_os.h"
#include "ola_log.h"
#include "ola_sdk_api.h"
#include "securec.h"
#include "mbedtls/base64.h"
#include "mbedtls/sha256.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/entropy.h"
#include "mbedtls/config.h"
#include "mbedtls/md.h"
#include "mbedtls/entropy_poll.h"
#include "mbedtls/gcm.h"

#define BASE64_MAX_LEN  4096
#define BITS_PER_BYTES  8

static int AesGcmParamCheck(const AesGcmParam *input, bool isDecript)
{
    if (input == NULL) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (input->key.data == NULL || ((input->key.len != OLA_AES_128_GCM_KEY_LEN) &&
        (input->key.len != OLA_AES_192_GCM_KEY_LEN) && (input->key.len != OLA_AES_256_GCM_KEY_LEN))) {
        OlaLogWarning("invalid key");
        return OLA_SECURE_ERR_AES_GCM_INVALID_KEY;
    }

    if ((input->iv.data == NULL) || (input->iv.len == 0)) {
        OlaLogWarning("invalid iv");
        return OLA_SECURE_ERR_AES_GCM_INVALID_IV;
    }

    if ((input->add.data == NULL) && (input->add.len != 0)) {
        OlaLogWarning("invalid add");
        return OLA_SECURE_ERR_AES_GCM_INVALID_ADD;
    }

    if ((input->inData.data == NULL) || (input->inData.len == 0)) {
        OlaLogWarning("invalid data");
        return OLA_SECURE_ERR_AES_GCM_INVALID_DATA;
    }

    if (isDecript) {
        if ((input->tag.len < OLA_AES_GCM_TAG_MIN_LEN) || (input->tag.len > OLA_AES_GCM_TAG_MAX_LEN)) {
            OlaLogWarning("invalid tag");
            return OLA_SECURE_ERR_AES_GCM_INVALID_TAG;
        }
    }

    return OLA_OK;
}

int OlaAesGcmEncript(const AesGcmParam *input, OlaBuffer *output, OlaBuffer *tag)
{
    int ret = AesGcmParamCheck(input, false);
    if (ret != OLA_OK) {
        OlaLogWarning("param check error");
        return ret;
    }
    if ((output == NULL) || (output->buf == NULL) || (tag == NULL) ||
        (input->inData.len > output->size) ||
        (tag->size < OLA_AES_GCM_TAG_MIN_LEN) || (tag->size > OLA_AES_GCM_TAG_MAX_LEN)) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    mbedtls_gcm_context context;
    mbedtls_gcm_init(&context);

    ret = mbedtls_gcm_setkey(&context, MBEDTLS_CIPHER_ID_AES, input->key.data, input->key.len * BITS_PER_BYTES);
    if (ret != 0) {
        OlaLogError("set key error [%d]\r\n", ret);
        mbedtls_gcm_free(&context);
        return OLA_SECURE_ERR_AES_GCM_SET_KEY;
    }

    ret = mbedtls_gcm_crypt_and_tag(&context, MBEDTLS_GCM_ENCRYPT, input->inData.len,
        input->iv.data, input->iv.len, input->add.data, input->add.len, input->inData.data,
        output->buf, tag->size, tag->buf);
    if (ret != 0) {
        OlaLogError("encrypt error [%d]\r\n", ret);
        mbedtls_gcm_free(&context);
        return OLA_SECURE_ERR_AES_GCM_ENCRIPT;
    }
    output->len = input->inData.len;
    tag->len = tag->size;

    mbedtls_gcm_free(&context);
    return OLA_OK;
}

int OlaAesGcmDecript(const AesGcmParam *input, OlaBuffer *output)
{
    int ret = AesGcmParamCheck(input, true);
    if (ret != OLA_OK) {
        OlaLogWarning("param check error");
        return ret;
    }
    if ((output == NULL) || (output->buf == NULL) || (input->inData.len > output->size)) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    mbedtls_gcm_context context;
    mbedtls_gcm_init(&context);
    ret = mbedtls_gcm_setkey(&context, MBEDTLS_CIPHER_ID_AES, input->key.data, input->key.len * BITS_PER_BYTES);
    if (ret != 0) {
        OlaLogError("set key error %d", ret);
        mbedtls_gcm_free(&context);
        return OLA_SECURE_ERR_AES_GCM_SET_KEY;
    }

    ret = mbedtls_gcm_auth_decrypt(&context, input->inData.len, input->iv.data, input->iv.len,
        input->add.data, input->add.len, input->tag.data, input->tag.len, input->inData.data, output->buf);
    if (ret != OLA_OK) {
        OlaLogError("decrypt error %d", ret);
        mbedtls_gcm_free(&context);
        return OLA_SECURE_ERR_AES_GCM_DECRIPT;
    }

    output->len = input->inData.len;
    mbedtls_gcm_free(&context);
    return OLA_OK;
}

int OlaEncodeBase64Data(const OlaData *inData, OlaData *outData)
{
    if ((inData == NULL) || (outData == NULL) || (inData->data == NULL) || (inData->len == 0) ||
        (outData->data != NULL) || (outData->len != 0)) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    size_t maxLen = 0;
    /* 获取base64编码所需buf长度 */
    int ret = mbedtls_base64_encode(NULL, 0, &maxLen, inData->data, inData->len);
    if ((ret != MBEDTLS_ERR_BASE64_BUFFER_TOO_SMALL) || (maxLen == 0)) {
        OlaLogError("get base64 max size error, ret[%d], maxLen[%u]", ret, maxLen);
        return OLA_MBEDTLS_ERR_GET_BASE64_SIZE;
    }
    if (maxLen > BASE64_MAX_LEN) {
        OlaLogWarning("base64 encode size too big");
        return OLA_COM_ERR_EXCEED_MAX_BUF_SIZE;
    }

    outData->data = (unsigned char *)OlaMalloc(maxLen);
    if (outData->data == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }
    (void)memset_s(outData->data, maxLen, 0, maxLen);
    ret = mbedtls_base64_encode(outData->data, maxLen, (size_t *)&outData->len, inData->data, inData->len);
    if ((ret != 0) || (outData->data[maxLen - 1] != 0)) {
        OlaLogError("base64 encode error, ret[%d]", ret);
        OlaPutNullAfterFree((void **)&outData->data);
        return OLA_MBEDTLS_ERR_BASE64_ENCODE;
    }
    return OLA_OK;
}

int OlaDecodeBase64Data(const OlaData *inData, OlaData *outData)
{
    if ((inData == NULL) || (outData == NULL) || (inData->data == NULL) || (inData->len == 0) ||
        (outData->data != NULL) || (outData->len != 0)) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    size_t outLen = 0;
    /* 获取base64编码所需buf长度 */
    int ret = mbedtls_base64_decode(NULL, 0, &outLen, inData->data, inData->len);
    if ((ret != MBEDTLS_ERR_BASE64_BUFFER_TOO_SMALL) || (outLen == 0)) {
        OlaLogError("get base64 output len error, ret[%d], outLen[%u]");
        return OLA_MBEDTLS_ERR_GET_BASE64_SIZE;
    }
    if (outLen > BASE64_MAX_LEN) {
        OlaLogWarning("base64 encode size too big");
        return OLA_COM_ERR_EXCEED_MAX_BUF_SIZE;
    }
    outData->data = (unsigned char *)OlaMalloc(outLen);
    if (outData->data == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }

    (void)memset_s(outData->data, outLen, 0, outLen);
    ret = mbedtls_base64_decode(outData->data, outLen, (size_t *)&outData->len, inData->data, inData->len);
    if (ret != 0) {
        OlaLogError("base64 decode error %d", ret);
        OlaPutNullAfterFree((void **)&outData->data);
        return OLA_MBEDTLS_ERR_BASE64_DECODE;
    }
    return OLA_OK;
}

int OlaComputeMdHmac(const OlaData *key, const OlaData *inData, unsigned char outBuf[OLA_MD_HMAC_LEN])
{
    if ((key == NULL) || (key->data == NULL) || (key->len == 0)) {
        OlaLogError("invalid key data");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    if ((inData == NULL) || (inData->data == NULL) || (inData->len == 0)) {
        OlaLogError("invalid in data");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    if (outBuf == NULL) {
        OlaLogError("invalid out data");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    int ret = mbedtls_md_hmac(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
        key->data, key->len, inData->data, inData->len, outBuf);
    if (ret != OLA_OK) {
        OlaLogError("comput md hmac failed. ret[%d]", ret);
        return OLA_MBEDTLS_ERR_COMPUTE_HMAC;
    }

    return OLA_OK;
}

/* sha256加密 */
int OlaSha256(const unsigned char *input, unsigned int inputLen, unsigned char output[OLA_SHA256_HEX_LEN])
{
    if (input == NULL || output == NULL || inputLen == 0) {
        OlaLogError("invalid in data");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    int ret = mbedtls_sha256_ret(input, inputLen, output, 0);
    if (ret != OLA_OK) {
        OlaLogError("sha256 failed. ret[%d]", ret);
        return OLA_MBEDTLS_ERR_SHA256;
    }

    return OLA_OK;
}